#pragma once

#include "compi/p2p.h"

#include <cstddef>
#include <string>
#include <thread>
#include <vector>

namespace compi
{

// 接收请求：基于固定类型与命名空间过滤（ANY_TAG），采用显式状态机
class RecvRequest
{
public:
    RecvRequest() = default;

    RecvRequest(Context & ctx, int src, void * payload, MPI_Datatype dt, TagNamespace ns) noexcept
        : _ctx(&ctx), _req(irecv_fixed(ctx, src, payload, dt, ns))
    {
        // 禁用底层自动处理/释放，仅保留解包，由上层按 Unpack→Process→Release 执行
        _req.set_strategy(StrategyUnpack);
        _state = RequestState::Pending;
    }

    // 通用模板构造：固定类型直接接收；非固定类型先接收字节并在处理阶段反序列化
    template <typename T>
    RecvRequest(Context & ctx, int src, T & payload, TagNamespace ns) noexcept : _ctx(&ctx)
    {
        if constexpr (types::is_fixed_type_v<T>)
        {
            _req = irecv_fixed(ctx, src, &payload, types::mpi_datatype_of<T>(), ns);
            _req.set_strategy(StrategyUnpack);
        }
        else if constexpr (std::is_same_v<T, std::string>)
        {
            _req = irecv_serialized(ctx, src, payload, ns);
            _req.set_strategy(StrategyUnpack);
        }
        else if constexpr (std::is_same_v<T, std::vector<std::byte>>)
        {
            _req = irecv_serialized(ctx, src, payload, ns);
            _req.set_strategy(StrategyUnpack);
        }
        else
        {
            _tmp_bytes.clear();
            _req = irecv_serialized(ctx, src, _tmp_bytes, ns);
            // 在底层请求上注册反序列化处理器，避免覆盖用户的 RecvRequest::_processor
            _req.set_processor([this, &payload](void *, const RequestStatus &)
                               { payload = deserialize<T>(_tmp_bytes); });
            _req.set_strategy(StrategyUnpack | StrategyProcess);
        }
        _state = RequestState::Pending;
    }

    // 启动请求（标记进入 Inflight）；具体发布由 poll() 推进
    void start() noexcept
    {
        if (_state == RequestState::Pending)
            _state = RequestState::Inflight;
    }

    // 推进一次进度；返回是否发生进展
    bool poll() noexcept
    {
        if (_state == RequestState::Completed || _state == RequestState::Canceled ||
            _state == RequestState::Failed)
            return false;
        bool progressed = _req.poll();
        // 若底层已完成，执行处理链：Unpack（已完成）→ Process → Release（除非处理异常）
        if (_req.completed())
        {
            auto st_res = _req.result();
            if (st_res.is_error())
            {
                _state = (_req.status().cancelled) ? RequestState::Canceled : RequestState::Failed;
                return true;
            }
            // 已解包（若为序列化类型）；执行处理
            try
            {
                if (_processor)
                {
                    _processor(_req.status());
                }
                // 正常处理后再显式释放
                (void)_req.release();
                _state = RequestState::Completed;
            }
            catch (...)
            {
                // Process 异常：不自动释放，交由外层处理避免悬空
                _state = RequestState::Failed;
            }
            return true;
        }
        return progressed;
    }

    // 阻塞等待完成，并按处理链执行；异常仅在 Process 阶段传播
    bool wait()
    {
        if (_state == RequestState::Completed)
            return true;
        try
        {
            const bool ok = _req.wait().is_success();
            auto st_res = _req.result();
            if (!ok || st_res.is_error())
            {
                _state = (_req.status().cancelled) ? RequestState::Canceled : RequestState::Failed;
                return false;
            }
            if (_processor)
                _processor(_req.status());
            (void)_req.release();
            _state = RequestState::Completed;
            return true;
        }
        catch (...)
        {
            // Process 异常：不自动释放
            _state = RequestState::Failed;
            return false;
        }
    }

    // 取消请求（最佳努力），并进入 Canceled
    bool cancel() noexcept
    {
        const bool ok = _req.cancel();
        if (ok)
            _state = RequestState::Canceled;
        return ok;
    }

    // 设置处理器（在完成时按状态机调用）
    void set_processor(std::function<void(const RequestStatus &)> processor) noexcept
    {
        _processor = std::move(processor);
    }

    RequestState state() const noexcept { return _state; }
    RequestStatus status() const noexcept { return _req.status(); }

private:
    Context * _ctx{nullptr};
    TransportRequest _req{};
    RequestState _state{RequestState::Pending};
    std::function<void(const RequestStatus &)> _processor{};
    std::vector<std::byte> _tmp_bytes{}; // 非固定类型接收的过渡字节缓冲
};

// 发送请求：基于固定类型命名空间标签，采用显式状态机
class SendRequest
{
public:
    SendRequest() = default;

    SendRequest(Context & ctx,
                int dest,
                const void * payload,
                MPI_Datatype dt,
                TagNamespace ns,
                SendMode mode = SendMode::Default) noexcept
        : _ctx(&ctx), _req(isend_fixed(ctx, dest, payload, dt, ns, mode))
    {
        // 发送侧无解包/处理，保持自动释放标签
        _req.set_strategy(StrategyRelease);
        _state = RequestState::Inflight; // 发送即刻进入 Inflight
    }

    // 序列化负载（字符串）构造
    SendRequest(Context & ctx,
                int dest,
                const std::string & payload,
                TagNamespace ns,
                SendMode mode = SendMode::Default) noexcept
        : _ctx(&ctx), _req(isend_serialized(ctx, dest, payload, ns, mode))
    {
        _req.set_strategy(StrategyRelease);
        _state = RequestState::Inflight;
    }

    // 序列化负载（字节向量）构造
    SendRequest(Context & ctx,
                int dest,
                const std::vector<std::byte> & payload,
                TagNamespace ns,
                SendMode mode = SendMode::Default) noexcept
        : _ctx(&ctx), _req(isend_serialized(ctx, dest, payload, ns, mode))
    {
        _req.set_strategy(StrategyRelease);
        _state = RequestState::Inflight;
    }

    // 通用模板构造：固定类型直接发送；非固定类型使用 serializer 编码发送
    template <typename T>
    SendRequest(Context & ctx,
                int dest,
                const T & payload,
                TagNamespace ns,
                SendMode mode = SendMode::Default) noexcept
        : _ctx(&ctx)
    {
        if constexpr (types::is_fixed_type_v<T>)
        {
            _req = isend_fixed(ctx, dest, &payload, types::mpi_datatype_of<T>(), ns, mode);
        }
        else if constexpr (std::is_same_v<T, std::string>)
        {
            _req = isend_serialized(ctx, dest, payload, ns, mode);
        }
        else if constexpr (std::is_same_v<T, std::vector<std::byte>>)
        {
            _req = isend_serialized(ctx, dest, payload, ns, mode);
        }
        else
        {
            _req = isend_serialized(ctx, dest, payload, ns, mode);
        }
        _req.set_strategy(StrategyRelease);
        _state = RequestState::Inflight;
    }

    bool poll() noexcept
    {
        if (_state == RequestState::Completed || _state == RequestState::Canceled ||
            _state == RequestState::Failed)
            return false;
        bool progressed = _req.poll();
        if (_req.completed())
        {
            auto st_res = _req.result();
            if (st_res.is_error())
            {
                _state = (_req.status().cancelled) ? RequestState::Canceled : RequestState::Failed;
            }
            else
            {
                _state = RequestState::Completed;
            }
            return true;
        }
        return progressed;
    }

    bool wait()
    {
        if (_state == RequestState::Completed)
            return true;
        const bool ok = _req.wait().is_success();
        auto st_res = _req.result();
        if (!ok || st_res.is_error())
        {
            _state = (_req.status().cancelled) ? RequestState::Canceled : RequestState::Failed;
            return false;
        }
        _state = RequestState::Completed;
        return true;
    }

    bool cancel() noexcept
    {
        const bool ok = _req.cancel();
        if (ok)
            _state = RequestState::Canceled;
        return ok;
    }

    RequestState state() const noexcept { return _state; }
    RequestStatus status() const noexcept { return _req.status(); }

private:
    Context * _ctx{nullptr};
    TransportRequest _req{};
    RequestState _state{RequestState::Pending};
};

// 进度引擎：提供 poll()/waitsome() 与 max_inflight 配额控制
class ProgressEngine
{
public:
    explicit ProgressEngine(std::size_t max_inflight = 0) noexcept : _max_inflight(max_inflight) {}

    void set_max_inflight(std::size_t m) noexcept { _max_inflight = m; }

    void add(RecvRequest req) { _recv_reqs.emplace_back(std::move(req)); }
    void add(SendRequest req) { _send_reqs.emplace_back(std::move(req)); }

    // 推进一次：先补齐配额（启动 Pending），再对 Inflight 进行 poll
    bool poll() noexcept
    {
        bool progressed = false;
        std::size_t inflight = inflight_count();
        // 启动 Pending 请求，尊重配额
        if (_max_inflight == 0 || inflight < _max_inflight)
        {
            const std::size_t credits =
                (_max_inflight == 0) ? SIZE_MAX : (_max_inflight - inflight);
            progressed = start_pending(credits) || progressed;
        }
        // 推进 Send/Recv
        for (auto & r : _recv_reqs)
            progressed = r.poll() || progressed;
        for (auto & s : _send_reqs)
            progressed = s.poll() || progressed;
        return progressed;
    }

    // 阻塞直到至少一个请求完成或失败/取消
    std::size_t waitsome()
    {
        std::size_t completed = 0;
        while (true)
        {
            (void)poll();
            for (auto & r : _recv_reqs)
            {
                if (r.state() == RequestState::Completed || r.state() == RequestState::Failed ||
                    r.state() == RequestState::Canceled)
                    ++completed;
            }
            for (auto & s : _send_reqs)
            {
                if (s.state() == RequestState::Completed || s.state() == RequestState::Failed ||
                    s.state() == RequestState::Canceled)
                    ++completed;
            }
            if (completed > 0)
                return completed;
            // 简单自旋让出调度
            std::this_thread::yield();
        }
    }

private:
    std::size_t inflight_count() const noexcept
    {
        std::size_t c = 0;
        for (const auto & r : _recv_reqs)
            c += (r.state() == RequestState::Inflight);
        for (const auto & s : _send_reqs)
            c += (s.state() == RequestState::Inflight);
        return c;
    }

    bool start_pending(std::size_t credits) noexcept
    {
        bool progressed = false;
        for (auto & r : _recv_reqs)
        {
            if (credits == 0)
                break;
            if (r.state() == RequestState::Pending)
            {
                r.start();
                --credits;
                progressed = true;
            }
        }
        for (auto & s : _send_reqs)
        {
            if (credits == 0)
                break;
            if (s.state() == RequestState::Pending)
            {
                // 发送请求在构造时已进入 Inflight，这里仅作为保险处理
                --credits;
                progressed = true;
            }
        }
        return progressed;
    }

private:
    std::size_t _max_inflight{0};
    std::vector<RecvRequest> _recv_reqs;
    std::vector<SendRequest> _send_reqs;
};
} // namespace compi